#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/1/26 16:06
@desc: 
"""
import cv2
import numpy as np
from skimage import transform


def convert_density_map(origin_map, shape=(320, 320), resize_shape=None):
    """
    将density map转化为其他shape

    Parameters
    ----------
    origin_map: 原始density map
    shape: 输出的density map尺寸
    resize_shape: 图像实际被resize后的大小

    Returns
    -------
    目标尺寸的density map
    """
    if resize_shape is None:
        resize_shape = shape
    else:
        assert resize_shape[0] <= shape[0] and resize_shape[1] <= shape[1], \
            'shape is smaller than resize shape !'
    pad_h, pad_w = (shape[0] - resize_shape[0]) // 2, (shape[1] - resize_shape[1]) // 2
    buffer_map = origin_map[pad_h:(pad_h + resize_shape[0]), pad_w:(pad_w + resize_shape[1])]
    summary = np.sum(buffer_map)
    if resize_shape[0] == shape[0] and resize_shape[1] == shape[1]:
        density_map = buffer_map
    else:
        density_map = transform.resize(buffer_map, shape, mode='symmetric', preserve_range=True)
    density_map = summary * density_map / np.sum(density_map)
    return density_map


def preprocess_image(image, mean=(80.197326, 75.26667, 72.822747),
                     shape=(256, 256), resize_shape=None):
    """
    对要检测的图像进行预处理， 包括resize、substract mean

    Parameters
    ----------
    image: 输入图像
    shape: 输出图像的shape
    mean: 各个channel的均值，依次为R、G、B
    resize_shape: 图像实际被resize后的大小

    Returns
    -------
    处理后的图像

    """
    if resize_shape is None:
        resize_shape = shape
    else:
        assert resize_shape[0] <= shape[0] and resize_shape[1] <= shape[1], \
            'shape is smaller than resize shape !'
    img = image.astype(np.float32)
    img = cv2.resize(img, resize_shape)
    img[:, :, 0] -= mean[0]
    img[:, :, 1] -= mean[1]
    img[:, :, 2] -= mean[2]
    pad_h, pad_w = (shape[0] - resize_shape[0]) // 2, (shape[1] - resize_shape[1]) // 2
    bias_h, bias_w = (shape[0] - resize_shape[0]) % 2, (shape[1] - resize_shape[1]) % 2
    img = np.pad(img, ((pad_h, pad_h + bias_h), (pad_w, pad_w + bias_w), (0, 0)),
                 'constant', constant_values=0)
    return img


def rotate_image(angle, image):
    """
    逆时针旋转一个图像 angle 度

    Parameters
    ----------
    angle: 要旋转的角度
    image: 原始图像

    Returns
    -------
    img_buf: 从旋转后的图像的中心切割出来与原始图像相同shape的图像
    offset: img_buf 相对于旋转后的图像的顶点（左上角）的起始坐标

    """
    shape = image.shape
    arc = angle * np.pi / 180
    post_size = (np.ceil(abs(shape[0] * np.cos(arc)) + abs(shape[1] * np.sin(arc))).astype(np.int32),
                 np.ceil(abs(shape[0] * np.sin(arc)) + abs(shape[1] * np.cos(arc))).astype(np.int32))
    offset = list(map(lambda x, y: (x - y) // 2, post_size, shape))
    if 0 == angle:
        img_buf = image.copy()
        img_buf = img_buf.astype(np.float32)
    else:
        img_buf = transform.rotate(image, angle, resize=True, preserve_range=True)
        img_buf = img_buf[offset[0]:(offset[0] + shape[0]), offset[1]:(offset[1] + shape[1]), :]
    return offset, img_buf


def get_revolve_matrix(angle, src_shape):
    """
    获得旋转矩阵 目的在于目标框回转到原来图像
    Parameters
    ----------
    angle：图像旋转角度
    src_shape：原图尺寸

    Returns：旋转矩阵
    -------

    """
    arc = angle * np.pi / 180
    post_shape = (np.ceil(abs(src_shape[0] * np.cos(arc)) + abs(src_shape[1] * np.sin(arc))).astype(np.int32),
                  np.ceil(abs(src_shape[0] * np.sin(arc)) + abs(src_shape[1] * np.cos(arc))).astype(np.int32))
    revolve_mat_1 = np.matrix([[1, 0, 0],
                               [0, -1, 0],
                               [-0.5 * src_shape[1], 0.5 * src_shape[0], 1]])
    revolve_mat_2 = np.matrix([[np.cos(arc), -np.sin(arc), 0],
                               [np.sin(arc), np.cos(arc), 0],
                               [0, 0, 1]])
    revolve_mat_3 = np.matrix([[1, 0, 0],
                               [0, -1, 0],
                               [0.5 * post_shape[1], 0.5 * post_shape[0], 1]])
    offset = list(map(lambda x, y: (x - y) // 2, post_shape, src_shape))
    return revolve_mat_1 * revolve_mat_2 * revolve_mat_3, post_shape, offset


def revolve_boxes(boxes, revolve_mat, src_shape, dst_shape=None):
    """
    旋转box回转到原图位置 正确标出人的位置
    revolve the boxes into the origin image
    Parameters
    ----------
    boxes：框出人的矩阵框
    revolve_mat：旋转矩阵
    src_shape：原图尺寸
    dst_shape：旋转后图片大小

    Returns
    -------

    """
    if dst_shape is None:
        dst_shape = src_shape
    result_boxes = []
    for box in boxes.tolist():
        points = np.matrix([[box[1] * src_shape[1], box[0] * src_shape[0], 1],
                            [box[3] * src_shape[1], box[0] * src_shape[0], 1],
                            [box[1] * src_shape[1], box[2] * src_shape[0], 1],
                            [box[3] * src_shape[1], box[2] * src_shape[0], 1]])
        revolved_points = points * revolve_mat
        min_x, min_y, max_x, max_y = dst_shape[1], dst_shape[0], 0, 0
        for point in revolved_points.tolist():
            min_x = min(min_x, point[0])
            min_y = min(min_y, point[1])
            max_x = max(max_x, point[0])
            max_y = max(max_y, point[1])
        result_boxes.append([min_y / dst_shape[0], min_x / dst_shape[1],
                             max_y / dst_shape[0], max_x / dst_shape[1]])
    return result_boxes
